SageMaker HyperPod のライフサイクルスクリプトと base-config の動きを理解する

SageMaker HyperPod のライフサイクルスクリプトと base-config の動きを理解する

SageMaker HyperPod を使う上で最初にハードルになりそうなライフサイクルスクリプトと base-config の動きをブログにしてみました。
Clock Icon2024.12.15

こんにちは!AWS 事業本部コンサルティング部のたかくに(@takakuni_)です。

SageMaker HyperPod では base-config というライフサイクルスクリプトが提供されています。

GitHub にも公開されており、とくにこだわりが無ければ、このスクリプトを使う。必要に応じて、特定のソフトウェアのインストールを付け足す等を行っていただくケースが多いのではないでしょうか。

https://github.com/aws-samples/awsome-distributed-training/tree/main/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/base-config

base-config を利用することでフルスクラッチで作り込むより、迅速に SageMaker HyperPod を使い始められます。

ただ、いくつかのファイルで構成されており、各ファイルがどのように構成されているのか、改めて理解したいと思いブログにしてみようと思います。

takakuni@ base-config % tree
.
├── add_users.sh
├── apply_hotfix.sh
├── config.py
├── hotfix
│   ├── hold-lustre-client.sh
│   └── mock-gpu-driver-deb.sh
├── lifecycle_script.py
├── mount_fsx.sh
├── on_create.sh
├── setup_mariadb_accounting.sh
├── setup_rds_accounting.sh
├── setup_sssd.py
├── shared_users_sample.txt
├── start_slurm.sh
└── utils
    ├── enroot.conf
    ├── fsx_ubuntu.sh
    ├── gen-keypair-ubuntu.sh
    ├── install_dcgm_exporter.sh
    ├── install_docker.sh
    ├── install_efa_node_exporter.sh
    ├── install_enroot_pyxis.sh
    ├── install_head_node_exporter.sh
    ├── install_prometheus.sh
    ├── install_slurm_exporter.sh
    ├── motd.sh
    ├── motd.txt
    ├── mount-s3.sh
    ├── pam_adopt_cgroup_wheel.sh
    ├── slurm_fix_plugstackconf.sh
    ├── ssh-to-compute.sh
    └── update_neuron_sdk.sh

3 directories, 30 files

重要度の高いファイル

ファイルがいくつかありますが、(個人的主観を含む)重要度の高いファイルは、たったの 3 つです。

  1. on_create.sh
  2. lifecycle_script.py
  3. provisioning_parameters.json

on_create.sh

on_create.sh は各インスタンスグループで利用する、ライフサイクルスクリプトのエントリポイントの役割を担います。

CreateCluster API の InstanceGroups の LifeCycleConfig.OnCreate で指定します。

つまり、各インスタンスからみれば on_create.sh から実行されていくイメージです。

CreateCluster.json
{
   "ClusterName": "string",
   "InstanceGroups": [
      {
         "ExecutionRole": "string",
         "InstanceCount": number,
         "InstanceGroupName": "string",
         "InstanceStorageConfigs": [
            { ... }
         ],
         "InstanceType": "string",
         "LifeCycleConfig": {
            "OnCreate": "string",
            "SourceS3Uri": "string"
         },
         "OnStartDeepHealthChecks": [ "string" ],
         "OverrideVpcConfig": {
            "SecurityGroupIds": [ "string" ],
            "Subnets": [ "string" ]
         },
         "ThreadsPerCore": number,
         "TrainingPlanArn": "string"
      }
   ],
   "NodeRecovery": "string",
   "Orchestrator": {
      "Eks": {
         "ClusterArn": "string"
      }
   },
   "Tags": [
      {
         "Key": "string",
         "Value": "string"
      }
   ],
   "VpcConfig": {
      "SecurityGroupIds": [ "string" ],
      "Subnets": [ "string" ]
   }
}

https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateCluster.html

以下は on_create.sh の中身にコメントを加えてみました。いくつか外部ファイルとやりとりが行われています。

ここで重要なポイントは PROVISIONING_PARAMETERS_PATH で指定した provisioning_parameters.jsonSAGEMAKER_RESOURCE_CONFIG_PATH で指定した /opt/ml/config/resource_config.jsonlifecycle_script.py の引数に渡されている部分です。

on_create.sh
#!/bin/bash

# コマンドの表示を有効にし、エラー時に終了するように設定
set -ex

# ログファイルのパスを定義し、必要なディレクトリを作成
LOG_FILE="/var/log/provision/provisioning.log"
mkdir -p "/var/log/provision"
touch $LOG_FILE

# コンソールとログファイルの両方にメッセージを出力する関数
logger() {
  echo "$@" | tee -a $LOG_FILE
}

# プロビジョニングパラメータファイルのパス
PROVISIONING_PARAMETERS_PATH="provisioning_parameters.json"

# SAGEMAKER_RESOURCE_CONFIG_PATH環境変数が設定されているか確認
if [[ -z "$SAGEMAKER_RESOURCE_CONFIG_PATH" ]]; then
  # 環境変数が設定されていない場合、デフォルトのパスを使用
  logger "Env var SAGEMAKER_RESOURCE_CONFIG_PATH is unset, trying to read from default location path"
  SAGEMAKER_RESOURCE_CONFIG_PATH="/opt/ml/config/resource_config.json"

  # デフォルトの設定ファイルが存在するか確認
  if [[ ! -f $SAGEMAKER_RESOURCE_CONFIG_PATH ]]; then
    logger "Env var SAGEMAKER_RESOURCE_CONFIG_PATH is unset and file does not exist: $SAGEMAKER_RESOURCE_CONFIG_PATH"
    logger "Assume vanilla cluster setup, no scripts to run. Exiting."
    exit 0
  fi
else
  # 環境変数が設定されている場合、ファイルの存在を確認
  logger "env var SAGEMAKER_RESOURCE_CONFIG_PATH is set to: $SAGEMAKER_RESOURCE_CONFIG_PATH"
  if [[ ! -f $SAGEMAKER_RESOURCE_CONFIG_PATH ]]; then
    logger "Env var SAGEMAKER_RESOURCE_CONFIG_PATH is set and file does not exist: $SAGEMAKER_RESOURCE_CONFIG_PATH"
    exit 1
  fi
fi

# 指定されたパラメータでライフサイクルスクリプトを実行
logger "Running lifecycle_script.py with resourceConfig: $SAGEMAKER_RESOURCE_CONFIG_PATH, provisioning_parameters: $PROVISIONING_PARAMETERS_PATH"

# Pythonスクリプトを実行し、標準出力と標準エラー出力の両方をログファイルに記録
python3 -u lifecycle_script.py \
  -rc $SAGEMAKER_RESOURCE_CONFIG_PATH \
  -pp $PROVISIONING_PARAMETERS_PATH >  >(tee -a $LOG_FILE) 2>&1

# Pythonスクリプトの終了コードを取得
exit_code=$?

# Pythonスクリプトと同じ終了コードで終了
exit $exit_code

lifecycle_script.py

lifecycle_script.py は各種ソフトウェアのインストールを行う Python スクリプトです。

引数で渡された provisioning_parameters.json をもとに、 SlurmNodeType クラスで各ノードタイプを決定し、インストールするソフトウェアを変更します。

lifecycle_script.py
#!/usr/bin/env python

import argparse
from enum import Enum
import json
import os
import socket
import subprocess
import sys
import time
from typing import Any, Dict, List, Optional, Tuple

from config import Config

# Slurmの設定ファイルのデフォルトパス
SLURM_CONF = os.getenv("SLURM_CONF", "/opt/slurm/etc/slurm.conf")

# Slurmクラスタ内のノードタイプを定義する列挙型
class SlurmNodeType(str, Enum):
    HEAD_NODE = "controller"    # コントローラーノード
    LOGIN_NODE = "login"       # ログインノード
    COMPUTE_NODE = "compute"   # 計算ノード

# Bashスクリプトを実行するためのユーティリティクラス
class ExecuteBashScript:
    def __init__(self, script_name: str):
        self.script_name = script_name

    def run(self, *args):
        print(f"Execute script: {self.script_name} {' '.join([str(x) for x in args])}")
        result = subprocess.run(["sudo", "bash", self.script_name, *args])
        result.check_returncode()
        print(f"Script {self.script_name} executed successully")

# SageMakerのリソース設定を管理するクラス
class ResourceConfig:
    INSTANCE_GROUP_NAME = "Name"
    INSTANCE_NAME = "InstanceName"
    CUSTOMER_IP_ADDRESS = "CustomerIpAddress"

    def __init__(self, path: str):
        with open(path, "r") as f:
            self._config = json.load(f)

    # IPアドレスからインスタンス情報を検索
    def find_instance_by_address(self, address) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
        for group in self._config["InstanceGroups"]:
            for instance in group["Instances"]:
                if instance.get(ResourceConfig.CUSTOMER_IP_ADDRESS) == address:
                    return group, instance
        return None, None

    # 指定されたグループ名に属するインスタンスのIPアドレスリストを取得
    def get_list_of_addresses(self, group_name) -> List[str]:
        for group in self._config["InstanceGroups"]:
            if group.get(ResourceConfig.INSTANCE_GROUP_NAME) != group_name:
                continue
            return [i.get(ResourceConfig.CUSTOMER_IP_ADDRESS) for i in group["Instances"]]
        return []

# プロビジョニングパラメータを管理するクラス
class ProvisioningParameters:
    WORKLOAD_MANAGER_KEY: str = "workload_manager"
    FSX_DNS_NAME: str = "fsx_dns_name"
    FSX_MOUNT_NAME: str = "fsx_mountname"

    def __init__(self, path: str):
        with open(path, "r") as f:
            self._params = json.load(f)

    @property
    def workload_manager(self) -> Optional[str]:
        return self._params.get(ProvisioningParameters.WORKLOAD_MANAGER_KEY)

    @property
    def fsx_settings(self) -> Tuple[str, str]:
        return self._params.get(ProvisioningParameters.FSX_DNS_NAME), self._params.get(ProvisioningParameters.FSX_MOUNT_NAME)

    @property
    def controller_group(self) -> Optional[str]:
        return self._params.get("controller_group")

    @property
    def login_group(self) -> Optional[str]:
        return self._params.get("login_group")

# 現在のノードのIPアドレスを取得する関数
def get_ip_address():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(('10.254.254.254', 1))
        IP = s.getsockname()[0]
    except Exception:
        IP = '127.0.0.1'
    finally:
        s.close()
    return IP

# Slurmの設定ファイルが利用可能になるまで待機
def wait_for_slurm_conf(controllers: List[str]) -> bool:
    sleep = 5 # sec
    timeout = 60  # sec
    for i in range(timeout // sleep):
        if not os.path.exists(SLURM_CONF):
            print("slurm.conf is not present. It is fine for login/compute nodes")
            return True
        with open(SLURM_CONF, "rt") as f:
            data = f.read()
            for ip in controllers:
                if ip in data:
                    print("slurm.conf found. It contains at least one controller address")
                    return True
        time.sleep(sleep)
    return False

# scontrolコマンドの出力を待機
def wait_for_scontrol():
    timeout = 120
    sleep = 5
    for i in range (timeout // sleep):
        try:
            output = subprocess.check_output(['scontrol', 'show', 'nodes'])
            if output.strip():
                print("Nodes registered with Slurm, Proceeding with install scripts.", output)
                return True
        except subprocess.CalledProcessError:
            pass
        print(f"Waiting for output. Retrying in {sleep} seconds...")
        time.sleep(sleep)
    print(f"Exceeded maximum wait time of {timeout} seconds. No output from scontrol.")
    return False

# メイン処理
def main(args):
    # 設定ファイルの読み込み
    params = ProvisioningParameters(args.provisioning_parameters)
    resource_config = ResourceConfig(args.resource_config)

    # FSxのマウント処理
    fsx_dns_name, fsx_mountname = params.fsx_settings
    if fsx_dns_name and fsx_mountname:
        print(f"Mount fsx: {fsx_dns_name}. Mount point: {fsx_mountname}")
        ExecuteBashScript("./mount_fsx.sh").run(fsx_dns_name, fsx_mountname, "/fsx")

    # ユーザー追加
    ExecuteBashScript("./add_users.sh").run()

    # Slurmクラスタの設定
    if params.workload_manager == "slurm":
        controllers = resource_config.get_list_of_addresses(params.controller_group)
        wait_for_slurm_conf(controllers)

        print("This is a slurm cluster. Do additional slurm setup")
        self_ip = get_ip_address()
        print(f"This node ip address is {self_ip}")

        # ノードタイプの判定
        group, instance = resource_config.find_instance_by_address(self_ip)
        if instance is None:
            raise ValueError("This instance not found in resource config. Can't process")
        print(group)

        # ノードタイプに応じた設定
        node_type = SlurmNodeType.COMPUTE_NODE
        if group.get("Name") == params.controller_group:
            node_type = SlurmNodeType.HEAD_NODE
        elif group.get("Name") == params.login_group:
            node_type = SlurmNodeType.LOGIN_NODE

        # ヘッドノードの場合はMariaDBのセットアップ
        if node_type == SlurmNodeType.HEAD_NODE:
            ExecuteBashScript("./setup_mariadb_accounting.sh").run()

        # 共通のセットアップスクリプトの実行
        ExecuteBashScript("./apply_hotfix.sh").run(node_type)
        ExecuteBashScript("./utils/motd.sh").run(node_type)
        ExecuteBashScript("./utils/fsx_ubuntu.sh").run()

        # Slurmの起動
        ExecuteBashScript("./start_slurm.sh").run(node_type, ",".join(controllers))

        # 監視機能のセットアップ(オプション)
        if Config.enable_observability:
            if node_type == SlurmNodeType.COMPUTE_NODE:
                ExecuteBashScript("./utils/install_docker.sh").run()
                ExecuteBashScript("./utils/install_dcgm_exporter.sh").run()
                ExecuteBashScript("./utils/install_efa_node_exporter.sh").run()

            if node_type == SlurmNodeType.HEAD_NODE:
                wait_for_scontrol()
                ExecuteBashScript("./utils/install_docker.sh").run()
                ExecuteBashScript("./utils/install_slurm_exporter.sh").run()
                ExecuteBashScript("./utils/install_head_node_exporter.sh").run()
                ExecuteBashScript("./utils/install_prometheus.sh").run()

        # Docker/Enroot/Pyxisのインストール(オプション)
        if Config.enable_docker_enroot_pyxis:
            ExecuteBashScript("./utils/install_docker.sh").run()
            ExecuteBashScript("./utils/install_enroot_pyxis.sh").run(node_type)

        # Neuron SDKのアップデート(オプション)
        if Config.enable_update_neuron_sdk:
            if node_type == SlurmNodeType.COMPUTE_NODE:
                ExecuteBashScript("./utils/update_neuron_sdk.sh").run()

        # SSSD(認証システム)のセットアップ(オプション)
        if Config.enable_sssd:
            subprocess.run(["python3", "-u", "setup_sssd.py", "--node-type", node_type], check=True)

        # SMHPの初期化(オプション)
        if Config.enable_initsmhp:
            ExecuteBashScript("./initsmhp.sh").run(node_type)

    print("[INFO]: Success: All provisioning scripts completed")

# スクリプトのエントリーポイント
if __name__ == "__main__":
    parser=argparse.ArgumentParser()
    parser.add_argument("-rc", "--resource_config", help="Resource config JSON file containing Ip_address of head, login and compute nodes")
    parser.add_argument("-pp", "--provisioning_parameters", help="Provisioning Parameters containing the head, login and compute ID/names")
    args=parser.parse_args()

    main(args)

config.py

Config クラスもみてみましょう。監視の有無等を設定していますね。イメージがどんどんついてきました。

config.py
# 基本設定パラメータ
class Config:

    # Docker/Enroot/Pyxisのインストール設定
    # デフォルトはTrue - コンテナ実行環境をインストール
    enable_docker_enroot_pyxis = True

    # 監視機能の有効化設定
    # Trueに設定すると以下がインストールされます:
    # - 計算ノード:DCGM ExporterとEFA Node Exporter
    # - コントローラーノード:Slurm ExporterとPrometheus
    enable_observability = False

    # PAMモジュールの設定
    # Trueに設定すると以下が有効になります:
    # - Slurmctldの再起動時の応答性問題を修正
    # - pam_slurm_adopt PAMモジュールをインストールし:
    #   - cgroupを使用してホストメモリ使用率を99% MaxRAMPercentに制限
    #   - ジョブが実行されていないノードへのSSHアクセスを防止
    enable_pam_slurm_adopt = False

    # Neuron SDKのアップデート設定
    # Trueに設定すると、計算ノード上のNeuron SDKを更新
    # (trnとinfクラスタのみに適用)
    enable_update_neuron_sdk = False

    # SSSD(System Security Services Daemon)の設定
    # Trueに設定すると、ActiveDirectory/LDAP統合のためのSSSDをインストール
    # 注:SssdConfigでの追加設定が必要
    enable_sssd = False

    # S3マウント機能の設定
    # Trueに設定すると、クラスタノードでS3をマウントポイントとして使用
    # - /mnt/<BucketName>にマウントするsystemctlサービスファイルが作成されます
    # - クラスタ実行ロールにS3権限の追加が必要です
    enable_mount_s3 = False

    # S3バケット名の設定
    # enable_mount_s3がTrueの場合は必須
    # 実際のデータバケット名を指定(例:"my-dataset-bucket")
    s3_bucket = ""

    # S3マウント設定の検証
    if enable_mount_s3 and not s3_bucket:
        raise ValueError("Error: A bucket name must be specified when enable_mount_s3 is True")


# ActiveDirectory/LDAP/SSSD設定パラメータ
class SssdConfig:

    # ドメイン名の設定
    # 不明な場合は"default"を使用可能
    domain = "default"

    # LDAPサーバーのURI(カンマ区切りで複数指定可能)
    ldap_uri = "ldaps://nlb-ds-xyzxyz.elb.us-west-2.amazonaws.com"

    # LDAPユーザー操作のベースDN
    ldap_search_base = "dc=hyperpod,dc=abc123,dc=com"

    # LDAP操作用のバインドDN
    ldap_default_bind_dn = "CN=ReadOnly,OU=Users,OU=hyperpod,DC=hyperpod,DC=abc123,DC=com"

    # 認証トークンタイプ
    # "password"または"obfuscated_password"(難読化パスワードを推奨)
    ldap_default_authtok_type = "obfuscated_password"

    # 認証トークン
    # 注:難読化されたパスワードを設定する必要があります(平文は不可)
    ldap_default_authtok = "placeholder"

    # SSH認証方式
    # "password"または"publickey"を指定
    ssh_auth_method = "publickey"

    # ホームディレクトリのパス
    # FSxを使用しないクラスタの場合は"/home/%u"に変更可能
    override_homedir = "/fsx/%u"

    # SSHログインを許可するグループ設定
    # ノードタイプごとにアクセス可能なグループを定義
    ssh_allow_groups = {
        "controller" : ["ClusterAdmin", "ubuntu"],
        "compute" : ["ClusterAdmin", "ClusterDev", "ubuntu"],
        "login" : ["ClusterAdmin", "ClusterDev", "ubuntu"],
    }

    # sudo権限を付与するグループ設定
    # ノードタイプごとにsudo権限を持つグループを定義
    sudoers_groups = {
        "controller" : ["ClusterAdmin", "ClusterDev"],
        "compute" : ["ClusterAdmin", "ClusterDev"],
        "login" : ["ClusterAdmin", "ClusterDev"],
    }

start_slurm.sh

start_slurm.sh では lifecycle_script.py の SlurmNodeType クラスで決定したノードタイプに応じて、 slurmctldslurmd の設定を切り替えています。

start_slurm.sh
#!/bin/bash

# このスクリプトはsudo権限で実行する必要があります
# 使用方法: start_slurm.sh <NODE_TYPE> [<CONTOLLER_ADDRESSES>]
# NODE_TYPEには以下のいずれかを指定:
# - controller: コントローラーノード
# - compute: 計算ノード
# - login: ログインノード

# エラー時即時終了とコマンドの表示を有効化
set -ex

# ログファイルのパスを設定
LOG_FILE="/var/log/provision/provisioning.log"

# コントローラーノードのIPアドレスを配列として取得
CONTROLLER_IP_VALUES=($2)

main() {
  echo "[INFO] START: Starting Slurm daemons"

  # ノードタイプに応じて適切なSlurmデーモンを起動
  if [[ $1 == "controller" ]]; then
    echo "[INFO] This is a Controller node. Start slurm controller daemon..."

    # コントローラーノードではslurmctldを有効化して起動
    systemctl enable --now slurmctld

    # コントローラーノードではslurmdを無効化
    # slurmdサービスファイルの名前を変更して起動できないようにする
    mv /etc/systemd/system/slurmd{,_DO_NOT_START_ON_CONTROLLER}.service \
        || { echo "Failed to mask slurmd, perhaps the AMI already masked it?" ; }

  elif [[ $1 == "compute" ]] || [[ $1 == "login" ]]; then
    echo "[INFO] Running on $1 node. Start slurm daemon..."

    # 注意:ログインノードでもslurmdを再起動する必要がある
    # これは/var/spool/slurmd/にslurm.confを取得するため
    # ただし、slurm.confにログインノードが含まれていないため、実際には実行されない

    # slurmdサービスの設定を更新
    # コントローラーのIPアドレスを設定してサービスファイルを生成
    SLURMD_OPTIONS="--conf-server $CONTROLLER_IP_VALUES" envsubst < /etc/systemd/system/slurmd.service > slurmd.service
    mv slurmd.service /etc/systemd/system/

    # systemdの設定を再読み込みしてslurmdを有効化・起動
    systemctl daemon-reload
    systemctl enable --now slurmd

    # 計算/ログインノードではslurmctldを無効化
    # slurmctldサービスファイルの名前を変更して起動できないようにする
    mv /etc/systemd/system/slurmctld{,_DO_NOT_START_ON_CONTROLLER}.service \
        || { echo "Failed to mask slurmctldd, perhaps the AMI already masked it?" ; }
  fi

  echo "[INFO] Start Slurm Script completed"
}

# スクリプトの実行
# 引数をすべてmain関数に渡す
main "$@"

provisioning_parameters.json

最後に provisioning_parameters.json です。このファイルは base-config には含まれていないため、自身で作成する必要があります。

何か記法がまとまったドキュメントがあるわけではないのですが、以下のような設定ができます。

provisioning_parameters.json
{
  "version": "1.0.0",
  "workload_manager": "slurm",
  "controller_group": "controller-machine",
  "login_group": "login-group",
  "worker_groups": [
    {
      "instance_group_name": "worker-group-1",
      "partition_name": ${instance_type}
    }
  ],
  "fsx_dns_name": "${FSX_ID}.fsx.${AWS_REGION}.amazonaws.com",
  "fsx_mountname": "${FSX_MOUNTNAME}"
}

参考:SageMaker HyperPod WorkShop

たとえば、ヘッドノードが 1 台。ワーカーノードが 2 台、ファイルシステム無しであれば、次のような設定値になります。

provisioning_parameters.json
{
  "version": "1.0.0",
  "workload_manager": "slurm",
  "controller_group": "controller-machine",
  "worker_groups": [
    {
      "instance_group_name": "worker-group",
      "partition_name": "ml.t3.medium"
    }
  ]
}

CreateCluster で渡すパラメーター例は以下になります。

重要なのは controller_group や instance_group_name、インスタンスタイプの名前が一致していることです。

CreateCluster.json
{
    "ClusterName": "ml-cluster",
    "InstanceGroups": [
      {
        "InstanceGroupName": "controller-machine",
        "InstanceType": "ml.t3.medium",
        "InstanceStorageConfigs": [
          {
            "EbsVolumeConfig": {
              "VolumeSizeInGB": 20
            }
          }
        ],
        "InstanceCount": 1,
        "LifeCycleConfig": {
          "SourceS3Uri": "s3://${BUCKET}/src",
          "OnCreate": "on_create.sh"
        },
        "ExecutionRole": "${ROLE}"
      },
      {
        "InstanceGroupName": "worker-group",
        "InstanceType": "ml.t3.medium",
        "InstanceCount": 2,
        "InstanceStorageConfigs": [
          {
            "EbsVolumeConfig": {
              "VolumeSizeInGB": 100
            }
          }
        ],
        "LifeCycleConfig": {
          "SourceS3Uri": "s3://${BUCKET}/src",
          "OnCreate": "on_create.sh"
        },
        "ExecutionRole": "${ROLE}"
      }
    ],
    "VpcConfig": {
      "SecurityGroupIds": ["$SECURITY_GROUP"],
      "Subnets":["$SUBNET_ID"]
    }
}

ライフサイクルスクリプトの相関関係

最後にライフサイクルスクリプトの相関関係について触れていきます。

イメージは次のとおりです。

Untitled(102).png

まず初めに、インスタンスグループに所属したインスタンスの tmp 領域に指定したバケット名、ディレクトリ名のスクリプトが保管されます。

続いて CreateCluster API の InstanceGroups の LifeCycleConfig.OnCreate で指定したファイルをもとにライフサイクルスクリプトが実行されます。今回ですと on_create.sh になります。

on_create.sh では先ほど触れたとおり、 lifecycle_script.py に対して provisioning_parameters.json/opt/ml/config/resource_config.json を引数に渡して実行しています。

lifecycle_script.py では引数で受け取った provisioning_parameters.json をもとに各ノードの役割を識別し、インストールするソフトウェアを切り替えている仕組みになります。

tmp 領域で実行されているかどうか、試しに on_create.shpwd コマンドを含めて、実行されているディレクトリを確認してみました。

on_create.sh
#!/bin/bash

set -ex

LOG_FILE="/var/log/provision/provisioning.log"
mkdir -p "/var/log/provision"
touch $LOG_FILE

# Function to log messages
logger() {
  echo "$@" | tee -a $LOG_FILE
}

PROVISIONING_PARAMETERS_PATH="provisioning_parameters.json"

# Log current directory and its contents
+ logger "Current working directory:"
+ pwd | tee -a $LOG_FILE
+ logger "Directory contents:"
+ ls -la | tee -a $LOG_FILE

if [[ -z "$SAGEMAKER_RESOURCE_CONFIG_PATH" ]]; then
  logger "Env var SAGEMAKER_RESOURCE_CONFIG_PATH is unset, trying to read from default location path"
  SAGEMAKER_RESOURCE_CONFIG_PATH="/opt/ml/config/resource_config.json"

  if [[ ! -f $SAGEMAKER_RESOURCE_CONFIG_PATH ]]; then
    logger "Env var SAGEMAKER_RESOURCE_CONFIG_PATH is unset and file does not exist: $SAGEMAKER_RESOURCE_CONFIG_PATH"
    logger "Assume vanilla cluster setup, no scripts to run. Exiting."
    exit 0
  fi
else
  logger "env var SAGEMAKER_RESOURCE_CONFIG_PATH is set to: $SAGEMAKER_RESOURCE_CONFIG_PATH"
  if [[ ! -f $SAGEMAKER_RESOURCE_CONFIG_PATH ]]; then
    logger "Env var SAGEMAKER_RESOURCE_CONFIG_PATH is set and file does not exist: $SAGEMAKER_RESOURCE_CONFIG_PATH"
    exit 1
  fi
fi

logger "Running lifecycle_script.py with resourceConfig: $SAGEMAKER_RESOURCE_CONFIG_PATH, provisioning_parameters: $PROVISIONING_PARAMETERS_PATH"

python3 -u lifecycle_script.py \
  -rc $SAGEMAKER_RESOURCE_CONFIG_PATH \
  -pp $PROVISIONING_PARAMETERS_PATH >  >(tee -a $LOG_FILE) 2>&1

exit_code=$?

exit $exit_code

ノードの tmp 領域で実行されていること、一連のファイル群がコピーされていることがわかります。

Current working directory:
/tmp/sagemaker-hyperpod-lifecycle-XXXXXXXXXXXX/config
Directory contents:
total 88
drwxr-xr-x 5 root root 4096 Dec 15 05:16 .
drwxr-xr-x 3 root root 4096 Dec 15 05:16 ..
-rw-r--r-- 1 root root 1655 Dec 15 04:44 add_users.sh
-rw-r--r-- 1 root root 723 Dec 15 04:44 apply_hotfix.sh
-rw-r--r-- 1 root root 2336 Dec 15 04:44 config.py
drwxr-xr-x 2 root root 4096 Dec 15 05:16 hotfix
drwxr-xr-x 2 root root 4096 Dec 15 05:16 initsmhp
-rw-r--r-- 1 root root 1009 Dec 15 04:44 initsmhp.sh
-rw-r--r-- 1 root root 8395 Dec 15 04:44 lifecycle_script.py
-rw-r--r-- 1 root root 3226 Dec 15 04:44 mount_fsx.sh
-rw-r--r-- 1 root root 1492 Dec 15 05:06 on_create.sh
-rw-r--r-- 1 root root 176 Dec 15 04:44 provisioning_parameters.json
-rw-r--r-- 1 root root 3577 Dec 15 04:44 setup_mariadb_accounting.sh
-rw-r--r-- 1 root root 2843 Dec 15 04:44 setup_rds_accounting.sh
-rw-r--r-- 1 root root 8262 Dec 15 04:44 setup_sssd.py
-rw-r--r-- 1 root root 88 Dec 15 04:44 shared_users_sample.txt
-rw-r--r-- 1 root root 1355 Dec 15 04:44 start_slurm.sh
drwxr-xr-x 2 root root 4096 Dec 15 05:16 utils
env var SAGEMAKER_RESOURCE_CONFIG_PATH is set to: /opt/ml/config/resource_config.json
Running lifecycle_script.py with resourceConfig: /opt/ml/config/resource_config.json, provisioning_parameters: provisioning_parameters.json
Execute script: ./add_users.sh

まとめ

以上、「SageMaker HyperPod のライフサイクルスクリプトと base-config の動きを理解する」でした。

SageMaker HyperPod を利用する上での最初のハードルになりそうな部分だったのでブログにしてみました。

もっと詳しく知りたい方は、以下の公式ドキュメントをご覧いただけますと幸いです。

https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/sagemaker-hyperpod-lifecycle-best-practices.html

AWS 事業本部コンサルティング部のたかくに(@takakuni_)でした!

Share this article

facebook logohatena logotwitter logo

© Classmethod, Inc. All rights reserved.